import pandas as pd
from app_config import get_engine_ts
from app_config import get_pro


# 定义一个函数来拼接每一组的ts_code
def split_and_concatenate(df, column_name, group_size):
    grouped = [df[i:i + group_size][column_name] for i in range(0, len(df), group_size)]
    concatenated_strings = [','.join(group) for group in grouped]
    return concatenated_strings


def found_amount():
    # 构建查询语句，使用参数化查询
    query = """
           SELECT DISTINCT ts_code, name, fund_type, invest_type, market, benchmark 
           FROM `fund_basic_e` 
           where benchmark is not null 
           and benchmark != '' 
           and delist_date is null
           """
    # 创建数据库引擎
    fund_all_e = pd.read_sql_query(query, get_engine_ts())

    # # 将结果保存到 Excel
    # fund_all_e.to_excel("fund_basic_e.xlsx", index=False)

    # 调用函数，按500个一组切分，并拼接ts_code
    result = split_and_concatenate(fund_all_e, 'ts_code', 500)

    share_list = []
    fund_nav_list = []
    dict_stock = {}

    # 打印结果
    for idx, concatenated_string in enumerate(result):
        print(f"Group {idx + 1}: {concatenated_string}")
        share_one = get_pro().fund_share(ts_code=concatenated_string, trade_date="20241205")
        share_list.append(share_one)

        fund_nav_one = get_pro().fund_nav(ts_code=concatenated_string, nav_date="20241205")
        fund_nav_list.append(fund_nav_one)

    share = pd.concat(share_list, axis=0)
    fund_nav = pd.concat(fund_nav_list, axis=0)

    fund_all_e = pd.merge(fund_all_e, share, on="ts_code", how="inner")
    fund_all_e = pd.merge(fund_all_e, fund_nav, on="ts_code", how="left")
    fund_all_e['total_amount'] = fund_all_e['unit_nav'] * fund_all_e['fd_share'] / 10000
    fund_all_e.sort_values(by="total_amount", ascending=False, inplace=True)
    fund_all_e.reset_index(inplace=True)
    fund_all_e.to_excel("fund_all_e.xlsx")
   
    # 沪深300
    mask = fund_all_e['benchmark'].isin(['沪深300指数', '沪深300指数收益率'])
    cs_index = fund_all_e[mask]
    cs_index.to_excel("沪深300.xlsx")
    amount__sum = cs_index['total_amount'].sum()
    dict_stock["沪深300"] = amount__sum
    fund_all_e = fund_all_e[~mask]
    # fund_all_e.to_excel("fund_basic_e.xlsx")

    # 上证50
    mask = fund_all_e['benchmark'].isin(['上证50指数', '上证50指数收益率'])
    cs_index = fund_all_e[mask]
    cs_index.to_excel("上证50.xlsx")
    dict_stock["上证50"] = cs_index['total_amount'].sum()
    fund_all_e = fund_all_e[~mask]
    # fund_all_e.to_excel("fund_basic_e.xlsx")

    # 中证500
    mask = fund_all_e['benchmark'].isin(['中证500指数', '中证500指数收益率'])
    cs_index = fund_all_e[mask]
    cs_index.to_excel("中证500.xlsx")
    dict_stock["中证500"] = cs_index['total_amount'].sum()
    fund_all_e = fund_all_e[~mask]
    # fund_all_e.to_excel("fund_basic_e.xlsx")

    # 科创板50
    mask = fund_all_e['benchmark'].isin(['上证科创板50成份指数', '上证科创板50成份指数收益率'])
    cs_index = fund_all_e[mask]
    cs_index.to_excel("科创板50.xlsx")
    dict_stock["创板50"] = cs_index['total_amount'].sum()
    fund_all_e = fund_all_e[~mask]
    fund_all_e.to_excel("fund_basic_e.xlsx")

    # 创业板
    mask = fund_all_e['benchmark'].isin(['创业板指数', '创业板指数收益率'])
    cs_index = fund_all_e[mask]
    cs_index.to_excel("创业板.xlsx")
    dict_stock["创业板"] = cs_index['total_amount'].sum()
    fund_all_e = fund_all_e[~mask]
    fund_all_e.to_excel("fund_basic_e.xlsx")

    # 中证1000
    mask = fund_all_e['benchmark'].isin(['中证1000指数', '中证1000指数收益率'])
    cs_index = fund_all_e[mask]
    cs_index.to_excel("中证1000.xlsx")
    dict_stock["中证1000"] = cs_index['total_amount'].sum()
    fund_all_e = fund_all_e[~mask]
    fund_all_e.to_excel("fund_basic_e.xlsx")

    # 创业板50
    mask = fund_all_e['benchmark'].isin(['创业板50指数', '创业板50指数收益率'])
    cs_index = fund_all_e[mask]
    cs_index.to_excel("创业板50.xlsx")
    dict_stock["创业板50"] = cs_index['total_amount'].sum()
    fund_all_e = fund_all_e[~mask]
    fund_all_e.to_excel("fund_basic_e.xlsx")

    # 中证A500
    mask = fund_all_e['benchmark'].isin(['中证A500指数', '中证A500指数收益率'])
    cs_index = fund_all_e[mask]
    cs_index.to_excel("中证A500.xlsx")
    dict_stock["中证A500"] = cs_index['total_amount'].sum()
    fund_all_e = fund_all_e[~mask]
    fund_all_e.to_excel("fund_basic_e.xlsx")

    # 上证180
    mask = fund_all_e['benchmark'].isin(['上证180指数', '上证180指数收益率'])
    cs_index = fund_all_e[mask]
    cs_index.to_excel("上证180.xlsx")
    dict_stock["上证180"] = cs_index['total_amount'].sum()
    fund_all_e = fund_all_e[~mask]
    fund_all_e.to_excel("fund_basic_e.xlsx")

    # 中证A50
    mask = fund_all_e['benchmark'].isin(['中证A50指数', '中证A50指数收益率'])
    cs_index = fund_all_e[mask]
    cs_index.to_excel("中证A50.xlsx")
    dict_stock["中证A50"] = cs_index['total_amount'].sum()
    fund_all_e = fund_all_e[~mask]
    fund_all_e.to_excel("fund_basic_e.xlsx")

    # 上证红利指数
    mask = fund_all_e['benchmark'].isin(['上证红利指数', '上证红利指数收益率'])
    cs_index = fund_all_e[mask]
    cs_index.to_excel("上证红利.xlsx")
    dict_stock["上证红利000015"] = cs_index['total_amount'].sum()
    fund_all_e = fund_all_e[~mask]
    fund_all_e.to_excel("fund_basic_e.xlsx")

    # 中证红利
    mask = fund_all_e['benchmark'].isin(['中证红利指数', '中证红利指数收益率'])
    cs_index = fund_all_e[mask]
    cs_index.to_excel("中证红利.xlsx")
    dict_stock["中证红利000922"] = cs_index['total_amount'].sum()
    fund_all_e = fund_all_e[~mask]
    fund_all_e.to_excel("fund_basic_e.xlsx")

    # 中证红利低波动
    mask = fund_all_e['benchmark'].isin(['中证红利低波动指数', '中证红利低波动指数收益率'])
    cs_index = fund_all_e[mask]
    cs_index.to_excel("中证红利低波动.xlsx")
    dict_stock["中证红利低波动"] = cs_index['total_amount'].sum()
    fund_all_e = fund_all_e[~mask]
    fund_all_e.to_excel("fund_basic_e.xlsx")

    # 中证红利低波动100
    mask = fund_all_e['benchmark'].isin(['中证红利低波动100指数', '中证红利低波动100指数收益率'])
    cs_index = fund_all_e[mask]
    cs_index.to_excel("中证红利低波动100.xlsx")
    dict_stock["中证红利低波动100"] = cs_index['total_amount'].sum()
    fund_all_e = fund_all_e[~mask]
    fund_all_e.to_excel("fund_basic_e.xlsx")

    print(str(dict_stock))

    # 将字典转换为列表，每个元素是一个包含键和值的字典
    index_amount = [{'key': k, 'value': v} for k, v in dict_stock.items()]

    # 将列表转换为 DataFrame
    df_index_amount = pd.DataFrame(index_amount)
    df_index_amount.to_excel("df_index_amount.xlsx")


if __name__ == '__main__':
    found_amount()
